import os
import pandas as pd
from pathlib import Path


class Parser:
    @staticmethod
    def _get_file_info(file_path):
        file_stats = os.stat(file_path)
        return {
            "filename": Path(file_path).name,
            "path": file_path,
            "time": file_stats.st_mtime,
            "size_b": file_stats.st_size,
        }

    @staticmethod
    def _get_folder_size(folder_path):
        total_size = 0
        for dirpath, dirnames, filenames in os.walk(folder_path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                total_size += os.path.getsize(fp)
        return total_size

    @staticmethod
    def _human_readable_size(size, decimal_places=2):
        for unit in ["B", "KB", "MB", "GB", "TB", "PB"]:
            if size < 1024.0 or unit == "PB":
                break
            size /= 1024.0
        return f"{size:.{decimal_places}f} {unit}"
    
    @staticmethod
    def _machine_readable_size(human_size:str):
        size, unit = human_size.split(" ")
        size = float(size)
        units = ["B", "KB", "MB", "GB", "TB", "PB"]
        return size * 1024 ** units.index(unit)

    @staticmethod
    def _human_readable_time(time):
        return pd.to_datetime(time, unit="s")

    @staticmethod
    def build_dataframe(directory, current_dir_only=True):
        # TODO improve performance by filtering out files that are not in the current directory
        data = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                file_path = os.path.join(root, file)
                data.append(Parser._get_file_info(file_path))
            for dir in dirs:
                dir_path = os.path.join(root, dir)
                data.append(
                    {
                        "filename": dir,
                        "path": dir_path,
                        "time": os.stat(dir_path).st_mtime,
                        "size_b": Parser._get_folder_size(dir_path),
                    }
                )
        if current_dir_only:
            data = [x for x in data if x["filename"] in os.listdir(directory)]

        df = pd.DataFrame(data)
        df["time"] = df["time"].apply(Parser._human_readable_time)
        df["size"] = df["size_b"].apply(Parser._human_readable_size)
        df = df[["filename", "path", "time", "size"]]
        return df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Parse directory and output file information.")
    parser.add_argument("-i", "--input", required=False, help="Input directory to parse")
    parser.add_argument("-o", "--output", required=False, help="Output CSV file path")
    args = parser.parse_args()

    target_path = os.path.expanduser(args.input) if args.input else input("Enter the target directory: ")
    df = Parser.build_dataframe(os.path.expanduser(target_path))
    df.sort_values(by="filename", ascending=True, inplace=True)
    os.makedirs("data", exist_ok=True)
    output_path = os.path.expanduser(args.output) if args.output else f"data/{target_path.rsplit("/")[-1]}.csv"
    df.to_csv(output_path, index=False)
